#include "toroidal.h"
#include <stdio.h>
#include "cuda_runtime.h"
#include "math.h"
#include "cudafunc.h"
#include <unistd.h>

field scalar(vfield B)
{
    double *data = new double[Nx*Ny];
    double *Br = B.getfr();
    double *Bz = B.getfz();
    double *Bp = B.getfphi();
    for (int i = 0; i < Nx*Ny; i++)
    {
        data[i] = sqrt(Br[i]*Br[i]+Bz[i]*Bz[i]+Bp[i]*Bp[i]);
    }
    return field(data);
}

void direction(vfield B, vfield dB)
{
    double *br = dB.getfr();
    double *bz = dB.getfz();
    double *bp = dB.getfphi();
    double *Br = B.getfr();
    double *Bz = B.getfz();
    double *Bp = B.getfphi();

    for (int i = 0; i < Nx*Ny; i++)
    {
        double length = sqrt(Br[i]*Br[i]+Bz[i]*Bz[i]+Bp[i]*Bp[i]);
        if (length == 0.0)
            continue;
        br[i] = Br[i]/length;
        bz[i] = Bz[i]/length;
        bp[i] = Bp[i]/length;
    }
    return;
}

field Apsi;
field Bsca;
double particle::q = 1.0;
double particle::m = 1.0;

int main(void)
{  
    Apsi = field();
    field Br = field();
    field Bz = field();
    field Bphi = field();
    Br.load_fromfile("Br.csv");
    Bz.load_fromfile("Bz.csv");
    Bphi.load_fromfile("Bphi.csv");
    Apsi.load_fromfile("Apsi.csv");
    // Apsi.vtkoutput("/home/kun/Desktop/new/psi.vtk");

    vfield B0 = vfield(&Br,&Bphi,&Bz);
    B0.scale(1.0/rhoh_vs_a);
    Bsca = scalar(B0);
    vfield b0 = vfield();
    direction(B0,b0);
    vfield Bg = Bsca.gradient();
    vfield curlb = vfield();
    cylv p0 = {4.6,0.0,0.4};
    b0.curl(&curlb);
    
    field aq = field();
    fiq A = fiq(&Bsca,&Bsca,4);
    fiq Aphi = fiq(4);
    vfiq gradA = vfiq(4);
    particle pa = particle(p0,0.6*rhoh_vs_a,0.5);

    // print(BAvp.interp(p0));
    // print(rhoh_vs_a*p0.r*B0.interp(p0));
    fiq phitest = fiq(6);
    perturb_hermitefunc(&phitest, 0, 4, 0.3, 20.0, true);

    // particle pa = particle({3.184033,0.0,0.0},0.034361,0.140607);
    vfiq E = vfiq(3);
    //perpgrad(phies, E);
    cylv* trace = new cylv[10000];
    world wld = world();
    wld.B0 = &B0;
    wld.cb = &curlb;
    wld.gB = &Bg;
    wld.E0 = &E;
    wld.Bsca = &Bsca;
    wld.b = &b0;

    for(int i = 0; i < 1; i++)
    {
        double vpar = pa.getvpara();
        cylv ccp = pa.getcrd();
        vec cp_xyz = cyl2xyz(ccp);
        pfvec em1 = wld.emfield(ccp);
        phsp dX_phsp1 = wld.push_rk2(em1, &pa, 0.01, vpar);
        cylv p1 = xyz2cyl(cp_xyz + dX_phsp1.X);
        pfvec em2 = wld.emfield(p1);
        phsp dX_phsp2 = wld.push_rk2(em2, &pa, 0.01, vpar+dX_phsp1.vp);
        cylv crdn = xyz2cyl(cp_xyz + 0.5*dX_phsp1.X + 0.5*dX_phsp2.X);
        pa.setcrd(crdn, vpar+0.5*dX_phsp1.vp+0.5*dX_phsp2.vp);
        //pa.Pphi();
        double Bnorm = sqrt(dot(B0.interp(crdn),B0.interp(crdn)));

        if(i%1500 == 0)
            printf("%lf\n", pa.Pphi());
        trace[i] = pa.getcrd();
        if(!checkboundary(pa.getcrd()))
            printf("\noutside!!");
    }
    // printf("\n");
    // print(cyl2xyz(pa.getcrd()));
    // printf("\n%f", pa.getvpara());
    // trace_vtkoutput("/home/kun/Desktop/new/trace.vtk",trace, 10000);
    // trace_output("/home/kun/Desktop/new/trace1.csv",trace, 3000);
   

    int Npart = 65536;
    ptcs ptcstest = ptcs(Npart);
    ptcstest.init_realistic(Bsca);

    fiq phies = fiq(3);
    perturb_hermitefunc(&phies , 0, 4, 0.3, 40.0, true);
    
    phies.vtkoutput("/home/kun/Desktop/perturb.vtk");
    wld.phies = &phies;
    wld.fastptcs = &ptcstest;

    double* Fdata = new double[Nx*Ny];
    for (int i = 0; i < Nx*Ny; i++)
    {
        int ix = i/Ny;
        int iy = i%Ny;
        Fdata[i] = ptcstest.pol_dstrb({Rmin+dr*ix,0.0,Zmin+dz*iy});
    }
    field Fpol = field(Fdata);
    wld.fast_distrib = &Fpol;
    /* debug of cuda kernel */ /*
    particle pb = particle({3.9,0.0,0.3},0.01,0.4);
    wld.kcp_single(pb, 0.1);
    for(int i = 0 ; i<Npart; i++)
    {
        wld.fastptcs->ppts[i].setcrd({3.9,0.0,0.3},0.4);
        wld.fastptcs->ppts[i].setmu(0.01);
    } */ 
    cupush_rk2(wld, 0.02, 512, 0.09, 3);
    // wld.wpi_multis_kc(0.1);
    wld.denf = new fiq[1];
    double* pt = new double[Nx*Ny];
    wld.fastptcs->scatter(pt);
    field(pt).vtkoutput("/home/kun/Desktop/kcp.vtk");
    wld.denf->vtkoutput("/home/kun/Desktop/kcpn.vtk");
    vfield E0 = vfield();
    ptcstest.pB = &Bsca;
    ptcstest.output_poloidal("/home/kun/Desktop/new/realsp.csv");
    //ptcstest.output_velocityspace("/home/kun/Desktop/new/realvl.csv");
    
    double prob = ptcstest.histo(0.4);
    printf("\nhisto:%lf\n", prob);

    wld.b = &b0;
    wld.Bsca = &Bsca;
    wld.vA = &Bsca;
    wld.init();

    /* debug of shear and flow operator*/
    /*
    vfiq B1perp = vfiq(4);
    vfiq uperp = vfiq(4);
    wld.shear(phies, B1perp);
    wld.flow(phies, uperp);
    field(B1perp.get_qfr()).vtkoutput("/home/kun/Desktop/new/B1.vtk");
    field(uperp.get_ifz()).vtkoutput("/home/kun/Desktop/new/u1.vtk");
    printf("inner product:%f\n",inner_vf(B1perp, B1perp));
    printf("flow energy:%f\n",inner_vf(uperp, uperp));
    */

    /* TAE test */
   
    /*
    int num_mode = 24;
    fiq* funcs = new fiq[num_mode];
    vfiq shears[num_mode];
    vfiq flows[num_mode];
    for(int i = 0; i<num_mode; i++)
    {
        bool iq = (i%2 == 0);
        funcs[i] = fiq(2);
        perturb_hermitefunc(funcs+i, i/12, (2+(i/2)%6), 0.36, 30.0, iq);
        shears[i] = vfiq(2);
        flows[i] = vfiq(2);
        wld.shear(funcs[i], shears[i]);
        wld.flow(funcs[i], flows[i]);
        printf("flb:%f\n",inner_vf(shears[i], shears[i]));
    }
    
    double K[num_mode*num_mode]; // stiffness matrix
    double M[num_mode*num_mode]; // mass matrix;
    for(int i = 0; i<num_mode*num_mode; i++)
    {
        K[i] = inner_vf(shears[i%num_mode], shears[i/num_mode]);
        M[i] = inner_vf(flows[i%num_mode], flows[i/num_mode]);
        // printf("%1.4f %c",M[i], i%4 == 3?'\n':' ');
    }
    double coefs[num_mode];
    eigen_gen(K, M, num_mode, 0, coefs); 

    fiq eigenmode = fiq(2);
    linear_suppos(coefs, funcs, num_mode, &eigenmode);
    eigenmode.vtkoutput("/home/kun/Desktop/eigenmode.vtk");
    /* verify particle wave interaction implementation */
    /*
    wld.phies = &eigenm;
    particle pb = particle({3.78,0.0,0.1},0.1*rhoh_vs_a,0.1);
    wld.wpi_single(pb, 0.1);
    return 0; */
}

